import os

import numpy as np
from osgeo import gdal, osr

from EnviLdx import EnviLdx
from Timer import Timer
from utils import find_files, list_files, move_file

# Parameters
a0 = -204.1781
a1 = 2.5512
a2 = -0.0029
a3 = 2.3468
a4 = 0.1376
a5 = 39.9991
a6 = -170.1487
epsilon = 1e-10  # Small constant to prevent division by zero


def save_tiff(output_path, data, geo_transform):
    driver = gdal.GetDriverByName('GTiff')
    if driver is None:
        raise RuntimeError("GTiff driver not available.")
    dataset = driver.Create(output_path, data.shape[1], data.shape[0], 1, gdal.GDT_Float32)
    if dataset is None:
        raise RuntimeError(f"Failed to create dataset at {output_path}. Check path and permissions.")
    dataset.GetRasterBand(1).WriteArray(data)
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)  # Assuming WGS84
    dataset.SetProjection(srs.ExportToWkt())
    dataset.SetGeoTransform(geo_transform)
    dataset = None
    pass


class Fy3dLd3(EnviLdx):

    def __init__(self, file_path):
        EnviLdx.__init__(self, file_path)
        self.ndvi_data = None
        self.lst_data = None
        self.tvdi_data = None

    def calculate_nvdi(self, band_data=None):
        if band_data is None:
            band_data = self.band_data

        ndvi_data = (band_data[:, :, 3].astype(np.float32) - band_data[:, :, 2].astype(np.float32)) / (
                band_data[:, :, 2].astype(np.float32) + band_data[:, :, 3].astype(np.float32) + epsilon)
        self.ndvi_data = ndvi_data
        return ndvi_data

    def calculate_lst(self, ndvi_data=None, band_data=None):
        if ndvi_data is None:
            ndvi_data = self.ndvi_data
        if band_data is None:
            band_data = self.band_data

        gvc = (ndvi_data - 0.2) / 0.3
        gvc[gvc < 0] = 0.0
        gvc[gvc > 1.0] = 1.0
        pv = gvc ** 2
        epsilon108 = 0.99 * pv + 0.9547 * (1 - pv)
        epsilon120 = 0.99 * pv + 0.9709 * (1 - pv)
        epsilon_comb = (epsilon108 + epsilon120) * 0.5
        d_epsilon = epsilon108 - epsilon120

        lst_data = np.zeros_like(band_data[:, :, 0], dtype=np.float32)
        for x in range(band_data.shape[0]):
            for y in range(band_data.shape[1]):
                if band_data[x, y, 1] > 200 and band_data[x, y, 23] < 3000:
                    lst_data[x, y] = -99
                elif ndvi_data[x, y] <= 0:
                    lst_data[x, y] = -98
                else:
                    delta_band = band_data[x, y, 23].astype(np.float32) - band_data[x, y, 24].astype(np.float32)
                    # Calculate intermediate terms
                    term1 = a0
                    term2 = a1 * band_data[x, y, 23] / 10.0
                    term3 = a2 * ((band_data[x, y, 23].astype(np.float32) + band_data[x, y, 24].astype(
                        np.float32)) / 20.0) ** 2
                    term4 = a3 * (
                            band_data[x, y, 23].astype(np.float32) - band_data[x, y, 24].astype(np.float32)) / 10.0
                    term5 = a4 * ((band_data[x, y, 23].astype(np.float32) - band_data[x, y, 24].astype(
                        np.float32)) / 10.0) ** 2
                    term6 = a5 * (1 - epsilon_comb[x, y])
                    term7 = a6 * d_epsilon[x, y]
                    term8 = -273.3

                    # Check intermediate terms for anomalies
                    lst_value = term1 + term2 + term3 + term4 + term5 + term6 + term7 + term8
                    lst_data[x, y] = lst_value
                    if band_data[x, y, 0] > 200 and band_data[x, y, 1] > 200 and band_data[x, y, 2] > 200:
                        lst_data[x, y] = 999

                    # Check final LST value for anomalies
                    if lst_data[x, y] > 1000 or lst_data[x, y] < -1000:
                        print(f"LST anomaly at ({x}, {y}): {lst_data[x, y]}")
                        print(f"term1: {term1}")
                        print(f"term2: {term2}")
                        print(f"term3: {term3}")
                        print(f"delta_band: {delta_band}")
                        print(f"term4: {term4}")
                        print(f"term5: {term5}")
                        print(f"term6: {term6}")
                        print(f"term7: {term7}")
                        print(f"term8: {term8}")
                        print(f"band23: {band_data[x, y, 23]}")
                        print(f"band24: {band_data[x, y, 24]}")
                        print(f"epsilon_comb: {epsilon_comb[x, y]}")
                        print(f"d_epsilon: {d_epsilon[x, y]}")
                        print(f"Final LST value: {lst_value}")

        self.lst_data = lst_data
        return lst_data

    def calculate_tvdi(self, ndvi_data=None, lst_data=None):
        if ndvi_data is None:
            ndvi_data = self.ndvi_data
        if lst_data is None:
            lst_data = self.lst_data

        ndvi_min = np.nanmin(ndvi_data)
        ndvi_max = np.nanmax(ndvi_data)

        # 不合理，就是以上值不参与TVDI计算
        indices = np.where((lst_data != -99) & (lst_data != -98) & (lst_data != 999))
        lst_min = np.nanmin(lst_data[indices])
        lst_max = np.nanmax(lst_data[indices])

        # Avoid divide by zero by adding a small epsilon
        if ndvi_max == ndvi_min:
            slope = 0
        else:
            slope = (lst_max - lst_min) / (ndvi_max - ndvi_min + epsilon)

        intercept = lst_min - slope * ndvi_min

        lst_pred = slope * ndvi_data + intercept
        tvdi_data = (lst_data - lst_min) / (lst_pred - lst_min + epsilon)

        # If lst[x, y] = -99    lst[x, y] = -98    lst[x, y] = 999,  TVDI[x, y]=-1
        indices = np.where((lst_data == -99) | (lst_data == -98) | (lst_data == 999))
        tvdi_data[indices] = -1

        self.tvdi_data = tvdi_data
        return tvdi_data

    def generate_three(self, ndvi_dir=None, lst_dir=None, tvdi_dir=None):
        directory, file_name = os.path.split(self.file_path)
        if not ndvi_dir:
            ndvi_dir = directory
        if not lst_dir:
            lst_dir = directory
        if not tvdi_dir:
            tvdi_dir = directory

        timer = Timer()
        timer.tick(f"开始处理 {file_name}")
        geo_transform = self.dataset.GetGeoTransform()

        ndvi_data = self.calculate_nvdi()
        ndvi_tiff_path = os.path.join(ndvi_dir, os.path.splitext(file_name)[0] + '_NDVI.TIF')
        save_tiff(ndvi_tiff_path, ndvi_data, geo_transform)

        timer.tick(f"{ndvi_tiff_path} 已生成。")
        lst_data = self.calculate_lst()
        lst_tiff_path = os.path.join(lst_dir, os.path.splitext(file_name)[0] + '_LST.TIF')
        save_tiff(lst_tiff_path, lst_data, geo_transform)
        timer.tick(f"{lst_tiff_path} 已生成。")

        tvdi_data = self.calculate_tvdi()
        tvdi_tiff_path = os.path.join(tvdi_dir, os.path.splitext(file_name)[0] + '_TVDI.TIF')
        save_tiff(tvdi_tiff_path, tvdi_data, geo_transform)
        timer.tick(f"{tvdi_tiff_path} 已生成。")

        timer.tick(f"{file_name} 处理完成。")
        pass


class Fy3dLd3Batch:

    def __init__(self, ld3_dir):
        # file_paths = find_files(ld3_dir, lambda f: f.upper().endswith('.LD3'))
        file_paths = list_files(ld3_dir, lambda f: f.upper().endswith('.LD3'))
        self.ld3s = []
        for file_path in file_paths:
            ld3 = Fy3dLd3(file_path)
            self.ld3s.append(ld3)
        pass

    def generate_three(self, ndvi_dir=None, lst_dir=None, tvdi_dir=None, ok_dir=None, error_dir=None):
        for ld3 in self.ld3s:
            try:
                ld3.generate_three(ndvi_dir, lst_dir, tvdi_dir)
                ld3.close()
                if ok_dir:
                    ok_file = move_file(ld3.file_path, ok_dir)
                    print(f"文件移动至完成目录：{ok_file}。")
            except ValueError as e:
                ld3.close()
                print(f"文处理失败：{e}")
                if error_dir:
                    error_file = move_file(ld3.file_path, error_dir)
                    print(f"文件移动至失败目录：{error_file}。")
            print("\n")
        pass
